import h5py
import numpy as np

# Open the HDF5 file
with h5py.File('data.h5', 'r') as file:
    # Assuming the dataset is stored under the name 'u_predict'
    data = file['u_predict'][:]

import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from mpl_toolkits.mplot3d import Axes3D

# Function to update the plot for each frame
def update_plot(frame_number, zs, plot):
    plt.cla()  # Clear the current axes
    for z in zs:  # Loop through each z-slice
        # Create a 2D array for the current z-slice at the given time frame
        data_slice = data[:, :, z, frame_number]
        # Create a contour plot at this z-level with semi-transparency
        plot.contourf(X, Y, data_slice, zdir='z', offset=z, alpha=0.5)

    # Setting the plot limits, titles, etc.
    plot.set_title('Time: {}'.format(frame_number))
    plot.set_xlim([0, data.shape[0]])
    plot.set_ylim([0, data.shape[1]])
    plot.set_zlim([0, len(zs)])

# Create a meshgrid for plotting
X, Y = np.meshgrid(np.arange(data.shape[0]), np.arange(data.shape[1]))
zs = range(data.shape[2])  # Range of z-slices

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

# Creating the animation
ani = FuncAnimation(fig, update_plot, frames=np.arange(data.shape[3]), fargs=(zs, ax), interval=100)

# To display the animation in a Jupyter notebook, use:
# from IPython.display import HTML
# HTML(ani.to_html5_video())

# To save the animation as an mp4 file
ani.save('animation.mp4', writer='ffmpeg')

plt.show()